# encoding:utf-8
from pymongo import UpdateOne
from zc_core.client.mongo_client import Mongo

mongo = Mongo()


# 过滤数据
def filter_data():
    update_bulk = list()
    src_list = mongo.list("data_20211031")
    for item in src_list:
        item_x = {
            '_id': item.get('_id'),
            'catalog1Id': item.get('catalog1Id'),
            'catalog1Name': item.get('catalog1Name'),
            'catalog2Id': item.get('catalog2Id'),
            'catalog2Name': item.get('catalog2Name'),
            'catalog3Id': item.get('catalog3Id'),
            'catalog3Name': item.get('catalog3Name'),
        }
        update_bulk.append(UpdateOne({'_id': item_x.get('_id')}, {'$set': item_x}, upsert=False))
    print(len(update_bulk))
    mongo.bulk_write('sku_20211101', update_bulk)
    mongo.close()
    print('任务完成~')


def fix_supplier_name():
    update_bulk = list()
    supplier_list = mongo.list("supplier_pool", fields={'_id': 1, 'name': 1})
    supplier_pool_dict = {}
    for data in supplier_list:
        supplier_pool_dict.update({data.get('_id'): data.get('name')})

    sku_list = mongo.list("sku_pool", fields={'_id': 1, 'supplierId': 1})
    sku_dict = {}
    for data in sku_list:
        sku_dict.update({data.get('_id'): data.get('supplierId')})
    full_list = mongo.list("data_20211113", fields={'_id': 1})
    for item in full_list:
        _id = item.get('_id')
        supplier_id = sku_dict.get(item.get('_id'))
        supplier_name = supplier_pool_dict.get(supplier_id)
        item = {
            '_id': item.get('_id'),
            'supplierId': supplier_id,
            'supplierName': supplier_name
        }
        update_bulk.append(UpdateOne({'_id': item.get('_id')}, {'$set': item}, upsert=False))
    print(len(update_bulk))
    mongo.bulk_write('data_20211113', update_bulk)
    mongo.close()
    print('任务完成~')


if __name__ == '__main__':
    fix_supplier_name()
